#!/usr/bin/env python3

import os
import subprocess
import sys

from argparse import ArgumentParser


def bitness_marker(target):
    return '(64bit)'


def is_elf(path):
    elf_magic = b'\x7fELF'
    with open(path, 'rb') as f:
        header = f.read(4)
        return header == elf_magic


def get_elf_interp(path, target):
    cmd = ['env', 'LC_ALL=C', f'{target}-readelf', '-l', path]
    out = subprocess.check_output(cmd, encoding='utf-8')
    return parse_elf_interp(out.split('\n'))


def get_required_sonames(path, target):
    cmd = ['env', 'LC_ALL=C', f'{target}-objdump', '-p', path]
    out = subprocess.check_output(cmd, encoding='utf-8')
    return parse_required_sonames(out.split('\n'))


OBJDUMP_SAMPLE = """
Dynamic Section:
  NEEDED               libreadline.so.7
  NEEDED               libdl.so.2
  NEEDED               libc.so.6
  INIT                 0x000000000001e000
  FINI                 0x00000000000a21e4
  INIT_ARRAY           0x00000000000c3df0
  INIT_ARRAYSZ         0x0000000000000008
  FINI_ARRAY           0x00000000000c3df8
  FINI_ARRAYSZ         0x0000000000000008
  GNU_HASH             0x0000000000000308
  STRTAB               0x000000000000deb8
  SYMTAB               0x0000000000003358
  STRSZ                0x0000000000006bd6
  SYMENT               0x0000000000000018
  DEBUG                0x0000000000000000
  PLTGOT               0x00000000000c56c8
  PLTRELSZ             0x00000000000018a8
  PLTREL               0x0000000000000007
  JMPREL               0x000000000001c6d8
  RELA                 0x0000000000015988
  RELASZ               0x0000000000006d50
  RELAENT              0x0000000000000018
  BIND_NOW             0x0000000000000000
  FLAGS_1              0x0000000008000001
  VERNEED              0x00000000000158d8
  VERNEEDNUM           0x0000000000000002
  VERSYM               0x0000000000014a8e
  RELACOUNT            0x0000000000000435

Version References:
  required from libdl.so.2:
    0x09691a75 0x00 09 GLIBC_2.2.5
  required from libc.so.6:
    0x06969191 0x00 10 GLIBC_2.11
    0x06969194 0x00 08 GLIBC_2.14
    0x0d696918 0x00 07 GLIBC_2.8
    0x06969195 0x00 06 GLIBC_2.15
    0x0d696914 0x00 05 GLIBC_2.4
    0x09691974 0x00 04 GLIBC_2.3.4
    0x0d696913 0x00 03 GLIBC_2.3
    0x09691a75 0x00 02 GLIBC_2.2.5

"""


def parse_required_sonames(out):
    versioned_refs_started = False
    dynamic_section_started = False

    current_soname = None
    needed = set()
    versioned = set()

    for line in out:
        orig_line = line
        line = line.strip()
        if dynamic_section_started:
            if line.startswith('NEEDED'):
                if versioned_refs_started:
                    raise ValueError('Unexpected NEEDED after Version References')
                _, soname = line.split()
                needed.add(soname)
                continue
            elif line == 'Version References:':
                versioned_refs_started = True
                dynamic_section_started = False
                continue
        elif versioned_refs_started:
            if line.startswith('required from'):
                try:
                    _, _, current_soname = line.split()
                    if current_soname.endswith(':'):
                        current_soname = current_soname[:-1]
                except:
                    raise ValueError('Invalid versioned reference: "%s"' % line)
            elif len(line) != 0:
                try:
                    addr, flag1, flag2, symver = line.split()
                except:
                    raise ValueError('Invalid versioned reference: "%s"' % line)
                if current_soname is None:
                    raise ValueError('Versioned reference to unknown SONAME')
                versioned.add((current_soname, symver))
                try:
                    needed.remove(current_soname)
                except KeyError:
                    pass
        else:
            if line == 'Dynamic Section:':
                dynamic_section_started = True

    return needed, versioned


READELF_SAMPLE = """
Elf file type is DYN (Shared object file)
Entry point 0x20860
There are 11 program headers, starting at offset 64

Program Headers:
  Type           Offset   VirtAddr           PhysAddr           FileSiz  MemSiz   Flg Align
  PHDR           0x000040 0x0000000000000040 0x0000000000000040 0x000268 0x000268 R   0x8
  INTERP         0x0002a8 0x00000000000002a8 0x00000000000002a8 0x00001c 0x00001c R   0x1
      [Requesting program interpreter: /lib64/ld-linux-x86-64.so.2]
  LOAD           0x000000 0x0000000000000000 0x0000000000000000 0x01df80 0x01df80 R   0x1000
  LOAD           0x01e000 0x000000000001e000 0x000000000001e000 0x0841ed 0x0841ed R E 0x1000
  LOAD           0x0a3000 0x00000000000a3000 0x00000000000a3000 0x01f070 0x01f070 R   0x1000
  LOAD           0x0c2df0 0x00000000000c3df0 0x00000000000c3df0 0x00537c 0x012d08 RW  0x1000
  DYNAMIC        0x0c44b8 0x00000000000c54b8 0x00000000000c54b8 0x000210 0x000210 RW  0x8
  NOTE           0x0002c4 0x00000000000002c4 0x00000000000002c4 0x000044 0x000044 R   0x4
  GNU_EH_FRAME   0x0ad700 0x00000000000ad700 0x00000000000ad700 0x002fa4 0x002fa4 R   0x4
  GNU_STACK      0x000000 0x0000000000000000 0x0000000000000000 0x000000 0x000000 RW  0x10
  GNU_RELRO      0x0c2df0 0x00000000000c3df0 0x00000000000c3df0 0x002210 0x002210 R   0x1

 Section to Segment mapping:
  Segment Sections...
   00     
   01     .interp 
   02     .interp .note.gnu.build-id .note.ABI-tag .gnu.hash .dynsym .dynstr .gnu.version .gnu.version_r .rela.dyn .rela.plt 
   03     .init .plt .plt.got .text .fini 
   04     .rodata .eh_frame_hdr .eh_frame 
   05     .init_array .fini_array .data.rel.ro .dynamic .got .data .bss 
   06     .dynamic 
   07     .note.gnu.build-id .note.ABI-tag 
   08     .eh_frame_hdr 
   09     
   10     .init_array .fini_array .data.rel.ro .dynamic .got 
"""


def parse_elf_interp(out):
    MARK = '[Requesting program interpreter:'
    for line in out:
        line = line.strip()
        if line.startswith(MARK):
            _, interp = line.split(':')
            interp = interp.strip()  
            if interp.endswith(']'):
                interp = interp[:-1]
            return interp


def test_parse_elf_interp():
    result = parse_elf_interp(READELF_SAMPLE.split('\n'))
    assert result == '/lib64/ld-linux-x86-64.so.2'


class XAutoRequires:
    def __init__(self, target, buildroot='/'):
        self._target = target
        self._needed = set()
        self._versioned_refs = set()
        self._interp = set()
        self._out = sys.stdout
        self._ready = False
        self._buildroot = buildroot 

    def _elf_files(self):
        for dirpath, dirnames, filenames in os.walk(self._buildroot):
            for name in filenames:
                path = os.path.join(dirpath, name)
                if is_elf(path):
                    yield path

    def get_requires(self):
        for file in self._elf_files():
            needed, versioned_refs = get_required_sonames(file, self._target)
            elf_interp = get_elf_interp(file, self._target)
            self._needed = self._needed.union(needed)
            self._versioned_refs = self._versioned_refs.union(versioned_refs)
            if elf_interp is not None:
                self._interp.add(elf_interp)

    def emit_requires(self):
        if not self._ready:
            self.get_requires()

        bitness = bitness_marker(self._target)
        for elf_interp in sorted(self._interp):
            print(f'Requires: {elf_interp}', file=self._out)
        for soname in sorted(self._needed):
            print(f'Requires: {soname}{bitness}', file=self._out)
        for soname, ver in sorted(self._versioned_refs):
            print(f'Requires: {soname}({ver}){bitness}', file=self._out)


def main():
    parser = ArgumentParser(description='Compute RPM Requires for ELF binaries')
    parser.add_argument('target', help='target architecture')
    parser.add_argument('buildroot', help='buildroot where files are installed')
    args = parser.parse_args()
    test_parse_elf_interp()
    areq = XAutoRequires(args.target, buildroot=args.buildroot)
    areq.emit_requires()


if __name__ == '__main__':
    main()
